package com.feriki.msync.process.source;

import com.alibaba.fastjson.JSONException;
import com.alibaba.fastjson.JSONObject;
import com.feriki.msync.meta.MilvusCollectionMeta;
import com.feriki.msync.constants.ConnectType;
import com.feriki.msync.process.sink.SinkHandler;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.apache.commons.lang3.tuple.Pair;

import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.List;
import java.util.stream.Stream;

/**
 * @author luozhida
 * @create 2024-07-10
 */
@Slf4j
public class JsonLineFileSourceHandler extends SourceHandler {

    private String sourceName;

    public JsonLineFileSourceHandler(String sourceName, long batchSize, long limit) {
        super(sourceName, ConnectType.JSONL, batchSize, limit);
        this.sourceName = sourceName;
    }

    @Override
    public Pair<Boolean, Long> readAndHandle(SinkHandler sinkHandler) {
        List<JSONObject> sinkDatas = new ArrayList<>();
        //文本里第二行数据才是需要迁移的数据，这里初始化设置为-1是为了后面遍历时与实际的数据对齐下标
        final long[] lineNumber = {1};
        Boolean isFirstLine = true;
        //成功的数据量
        final long[] successNumber = {0};
        try {
            //使用Files.lines读取大文件有性能优势
            Stream<String> fileStream = Files.lines(Paths.get(sourceName), StandardCharsets.UTF_8)
                    .filter(StringUtils::isNotBlank);
            if (getLimit() > 0) {
                //限制数据，+1的原因是首行是元数据，不是正式的数据
                fileStream = fileStream.limit(getLimit() + 1);
                log.info("即将对源文件【{}】执行处理,限制{}条",getName(),getLimit());
            }else{
                log.info("即将对源文件【{}】执行处理",getName());
            }
            fileStream.forEach(line -> {
                if (lineNumber[0] == 1) {
                    //首行是元数据
                    MilvusCollectionMeta meta;
                    try {
                        meta = JSONObject.parseObject(line, MilvusCollectionMeta.class);
                    } catch (JSONException e) {
                        log.error("文件{}首行是非合法的元数据,读取到的内容:{}", sourceName, line);
                        throw new InnerException("");
                    }
                    if (!sinkHandler.prepare(sourceName, meta)) {
                        throw new InnerException("");
                    }
                } else{
                    //处理正式数据
                    JSONObject lineJson = JSONObject.parseObject(line);
                    sinkDatas.add(lineJson);
                    if (sinkDatas.size() % getBatchSize() == 0) {
                        long success = sinkHandler.handler(sourceName, sinkDatas);
                        successNumber[0] = successNumber[0] + success;
                        sinkDatas.clear();
                    }
                }
                lineNumber[0] ++;
            });
            if (!sinkDatas.isEmpty()) {
                //对剩余的数据进行处理
                long success = sinkHandler.handler(sourceName, sinkDatas);
                successNumber[0] = successNumber[0] + success;
            }
            return Pair.of(true, successNumber[0]);
        } catch (IOException e) {
            log.error("读取文件{}出错", sourceName, e);
            return Pair.of(false, null);
        } catch (InnerException e) {
            return Pair.of(false, null);
        }
    }

    @Override
    public void destroy() {

    }


    class InnerException extends RuntimeException {
        public InnerException(String message) {
            super(message);
        }
    }
}
